'''
https://leetcode.cn/problems/best-time-to-buy-and-sell-stock-iv/description/
'''
from functools import cache
from typing import List


class Solution:

    def maxProfit(self, k: int, prices: List[int]) -> int:
        n = len(prices)

        # i: 第i个股票
        # rest: 还能做几笔交易
        @cache
        def f(i, rest):
            if rest == 0:
                return 0
            if i == n:
                return 0
            res = f(i + 1, rest)  # 当前不买也不卖
            for j in range(i + 1, n):
                res = max(res, prices[j] - prices[i] + f(j, rest - 1))
            return res

        return f(0, k)

    def maxProfit2(self, k: int, prices: List[int]) -> int:
        n = len(prices)
        dp = [[0] * (k + 1) for _ in range(n + 1)]
        # 第一维度依赖后边的，第二维度依赖前边的
        # 从下往上，从左往右
        for i in range(n - 1, -1, -1):
            for rest in range(1, k + 1):
                res = dp[i + 1][rest]
                for j in range(i + 1, n):
                    res = max(res, dp[j][rest - 1] + prices[j] - prices[i])
                dp[i][rest] = res
        print(dp)
        return dp[0][k]

    def maxProfit3(self, k: int, prices: List[int]) -> int:
        n = len(prices)
        dp = [[0] * (k + 1) for _ in range(n + 1)]
        # 第一维度依赖后边的，第二维度依赖前边的（从下往上，从左往右）
        # dp[i][rest]
        # 依赖 dp[i+1][rest],
        #   dp[n-1][rest-1]+ prices[n-1] - prices[i]
        #   dp[n-2][rest-1]+ prices[n-2] - prices[i]
        #   ...
        #   dp[i+2][rest-1]+ prices[i+2] - prices[i]
        #   dp[i+1][rest-1]+ prices[i+1] - prices[i]
        # 如果在推的过程中能记住max(dp[t][rest-1]+ prices[t])  t属于[i+1, n-1]
        # 则可以优化枚举，变为拿来直接用
        for rest in range(1, k + 1):
            best = dp[n-1][rest-1] + prices[n-1]
            for i in range(n - 2, -1, -1):
                res = dp[i + 1][rest]
                res = max(res, best - prices[i])
                dp[i][rest] = res
                best = max(best, dp[i][rest-1] + prices[i])
        return dp[0][k]

prices = [2,4,1]
k = 2
res = Solution().maxProfit2(k, prices)
res2 = Solution().maxProfit3(k, prices)
print(res, res2)
